import pandas as pd
import numpy as np
import random
import matplotlib.pyplot as plt
import os
import time
from concurrent.futures import ThreadPoolExecutor
from scipy.optimize import minimize

# 定义蚁群算法类
class AntColonyOptimization:
    def __init__(self, num_ants, num_iterations, alpha, beta, rho, q0, data, param_ranges):
        self.num_ants = num_ants
        self.num_iterations = num_iterations
        self.alpha = alpha
        self.beta = beta
        self.rho = rho
        self.q0 = q0
        self.data = data
        self.param_ranges = param_ranges
        self.best_solution = None
        self.best_error = float('inf')
        self.error_history = []

    def objective_function(self, params):
        k, a, b = params
        P_actual = self.data['P']
        f = self.data['f']
        B = self.data['B']
        P_predicted = k * (f ** a) * (B ** b)

        # 使用相对误差
        error = np.mean(((P_actual - P_predicted) / P_actual) ** 2)
        return error

    def update_pheromones(self, pheromones, solutions, errors):
        for i in range(len(solutions)):
            for j in range(len(solutions[i])):
                pheromones[j] += (self.q0 / (errors[i] + 1e-10))  # 避免除零
        pheromones *= (1 - self.rho)
        return pheromones

    def run(self):
        pheromones = np.ones(3)
        for iteration in range(self.num_iterations):
            solutions = []
            errors = []

            # 并行计算多个蚂蚁的解
            with ThreadPoolExecutor() as executor:
                futures = [executor.submit(self.random_solution) for _ in range(self.num_ants)]
                for future in futures:
                    params, error = future.result()
                    solutions.append(params)
                    errors.append(error)
                    if error < self.best_error:
                        self.best_error = error
                        self.best_solution = params

            pheromones = self.update_pheromones(pheromones, solutions, errors)
            self.error_history.append(self.best_error)

        return self.best_solution, self.best_error

    def random_solution(self):
        # 随机生成初始解
        initial_params = [random.uniform(*self.param_ranges['k']),
                          random.uniform(*self.param_ranges['a']),
                          random.uniform(*self.param_ranges['b'])]

        # 根据误差条件决定是否局部优化
        initial_error = self.objective_function(initial_params)
        if initial_error > 0.05:  # 对误差较大的解进行优化
            result = minimize(self.objective_function, initial_params, method='L-BFGS-B', bounds=[
                self.param_ranges['k'], self.param_ranges['a'], self.param_ranges['b']
            ])
            params = result.x
            error = result.fun
        else:
            params = initial_params
            error = initial_error
        return params, error

# 读取CSV文件路径
file_paths = [
    r'D:\2024华为杯数学建模\2024C题-数据驱动与电磁损耗\C题\附件一数据拆解\Q2\Q2磁芯材料1正弦波形.csv',
    r'D:\2024华为杯数学建模\2024C题-数据驱动与电磁损耗\C题\附件一数据拆解\Q2\Q2磁芯材料2正弦波形.csv',
    r'D:\2024华为杯数学建模\2024C题-数据驱动与电磁损耗\C题\附件一数据拆解\Q2\Q2磁芯材料3正弦波形.csv',
    r'D:\2024华为杯数学建模\2024C题-数据驱动与电磁损耗\C题\附件一数据拆解\Q2\Q2磁芯材料4正弦波形.csv'
]

# 合并所有数据
data_frames = []
for file_path in file_paths:
    if not os.path.exists(file_path):
        raise FileNotFoundError(f"文件不存在: {file_path}")
    try:
        df = pd.read_csv(file_path, encoding='gbk')
    except UnicodeDecodeError:
        df = pd.read_csv(file_path, encoding='ISO-8859-1')
    data_frames.append(df)

data = pd.concat(data_frames, ignore_index=True)
num_columns = data.shape[1]
new_column_names = [f'Column_{i+1}' for i in range(num_columns)]
data.columns = new_column_names
data.rename(columns={'Column_2': 'f', 'Column_3': 'P', 'Column_5': 'B'}, inplace=True)

# 定义参数范围
param_ranges = {
    'k': (0.1, 10),  # k的范围
    'a': (0, 3),     # a的范围
    'b': (0, 3)      # b的范围
}

# 记录运行时间
start_time = time.time()

# 实例化蚁群算法
aco = AntColonyOptimization(
    num_ants=20,  # 蚂蚁数量
    num_iterations=300,  # 迭代次数
    alpha=1.0,
    beta=5.0,
    rho=0.1,
    q0=0.9,
    data=data,
    param_ranges=param_ranges
)

# 运行蚁群算法
best_params, best_error = aco.run()
end_time = time.time()

# 输出运行时间
print(f"运行时间: {end_time - start_time:.2f} 秒")

# 输出最优参数和误差
k, a, b = best_params
print(f"拟合参数:\n k = {k}\n a = {a}\n b = {b}\n 最小误差 = {best_error}")

# 可视化实际值与预测值
f = data['f']
B = data['B']
P_actual = data['P']
P_predicted = k * (f ** a) * (B ** b)

plt.figure(figsize=(10, 6))
plt.scatter(P_actual, P_predicted, color='blue', label='pre_P')
plt.plot([P_actual.min(), P_actual.max()], [P_actual.min(), P_actual.max()], 'r--', label='y=x')
plt.xlabel('true P')
plt.ylabel('pre P')
plt.title('true P VS pre P')
plt.legend()
plt.grid()
output_plot_path = r'D:\2024华为杯数学建模\2024C题-数据驱动与电磁损耗\C_code\实际与预测对比.png'
plt.savefig(output_plot_path)
plt.show()

# 绘制迭代曲线
plt.figure(figsize=(10, 6))
plt.plot(range(1, len(aco.error_history) + 1), aco.error_history, label='best error')
plt.xlabel('num')
plt.ylabel('error')
plt.title('error evolution')
plt.legend()
plt.grid()
output_convergence_path = r'D:\2024华为杯数学建模\2024C题-数据驱动与电磁损耗\C_code\迭代曲线.png'
plt.savefig(output_convergence_path)
plt.show()

# 保存最优结果到CSV文件
output_csv_path = r'D:\2024华为杯数学建模\2024C题-数据驱动与电磁损耗\C_code\最优拟合结果.csv'
result_df = pd.DataFrame({
    '参数': ['k', 'a', 'b', '最小误差'],
    '值': [k, a, b, best_error]
})
result_df.to_csv(output_csv_path, index=False)
